//============================================================================
// Name        : cg.c
// Author      : Yugi
// Version     : 1.0
// Description : 
//============================================================================

#include "cg.h"

/* 
* Generate a new call gragh.
* 
* RETURN:	new cg
* NOTES:	Return value should be freed by user
*/
PCG cg_new(){
	PCG pcg = (PCG) malloc(sizeof(CG));
	pcg->basic_blocks = (app_pc *) malloc(DEFAULT_CG_SIZE * 4);
	pcg->head = NULL;
	pcg->tail = NULL;
	pcg->size = 0;
	pcg->multiple = 1;
	pcg->stack = NULL;
	return pcg;
}

/* 
* Insert a basic block to the call gragh.
* 
* REQUIRE:	pcg not NULL
*			pc in text seg
* RETURN:	Index of the target basic block, if the basic block
*			exist, return the existing basic block index
* MODIFY:	pcg
*/
unsigned int cg_add_bb(PCG pcg, app_pc pc){
	unsigned int index;

	/* Return the index of bb if existed */
	for (index = 0; index<pcg->size; index++)
		if (pcg->basic_blocks[index] == pc)
			return index;

	pcg->basic_blocks[pcg->size] = pc;
	pcg->size++;

	/* If pcg->size meets the max size, twice it */
	if (pcg->size == DEFAULT_CFG_SIZE * pcg->multiple){
		app_pc *basic_blocks = pcg->basic_blocks;

		pcg->multiple++;
		pcg->basic_blocks = (app_pc *) malloc(DEFAULT_CG_SIZE*4*pcg->multiple);

		for (index = 0; index<pcg->size; index++)
			pcg->basic_blocks[index] = basic_blocks[index];

		free(basic_blocks);
	}

	return pcg->size-1;
}

/* 
* Insert an edge to the call gragh.
* 
* REQUIRE:	pcg not NULL
*			src & dst distinguishes, less than pcg->size
* RETURN:	true if edge added
* MODIFY:	pcg
*/
bool cg_add_edge(PCG pcg, unsigned int src, unsigned int dst){
		/* 
	* Iterate through pcfg->edges, note that pcfg->edges is represented
	* by pcfg->head & pcfg->tail
	*/
	PEDGE edge = pcg->head;
	
	while (edge){
		if (edge->src==src && edge->dst==dst)
			return false;

		edge = edge->next;
	}
	
	/* Add to pcfg->edges */
	edge = (PEDGE) malloc(sizeof(EDGE));
	edge->src = src;
	edge->dst = dst;
	edge->next = NULL;

	if (pcg->head){
		pcg->tail->next = edge;
		pcg->tail = edge;
	}else
		pcg->head = pcg->tail = edge;

	return true;
}

/* 
* Construct a call graph. 
* 
* REQUIRE:	pcg not NULL
*			module not NULL
*			bb not NULL
* MODIFY:	pcg
*/
void cg_parse(void *drcontext, instrlist_t *bb, module_data_t *module, PCG pcg){
	app_pc pc;				// basic block address
	instr_t *pLast;			// Last instruction of basic block
	PINT_STACK pStack;		// Call stack pointer

	/* Initial */
	pc = instr_get_app_pc(instrlist_first(bb));
	pLast = instrlist_last(bb);

	/* Cares about app code */
	if (pc<(module->start) || pc>(module->end))
		return;

	/* Entry point of the app */
	if (pc == module->entry_point){
		pStack = (PINT_STACK) malloc(sizeof(INT_STACK));
		pStack->index = cg_add_bb(pcg, pc);
		pStack->previous = NULL;
		pcg->stack = pStack;
	}

	/* Empty call stack */
	if (!pcg->stack)
		return;

	switch (instr_get_opcode(pLast)){
		case OP_call:
			pc = instr_get_branch_target_pc(pLast);

			if (pc){
				unsigned int index = cg_add_bb(pcg, pc);
				cg_add_edge(pcg, pcg->stack->index, index);

				/* If it is new basic block, add to call stack */
				if (index == pcg->size-1){
					pStack = (PINT_STACK) malloc(sizeof(INT_STACK));
					pStack->index = index;
					pStack->previous = pcg->stack;
					pcg->stack = pStack;
				}
			}

			break;
		case OP_ret:
			/* Remove the top basic block on call stack */
			pStack = pcg->stack;
			pcg->stack = pcg->stack->previous;
			free(pStack);
			break;
		default:
			;
	}
}

/* 
* Print call gragh
* 
* REQUIRE:	pcg not NULL
*			outfile not NULL
*/
void cg_print(void *drcontext, file_t outfile, PCG pcg){
	PEDGE edge = pcg->head;

	dr_fprintf(outfile, "digraph CG {\n");

	/* Iterate through pcg->edges */	
	while (edge){
		dr_fprintf(outfile, "\"0x%.8x\" -> \"0x%.8x\"\n", 
			pcg->basic_blocks[edge->src], pcg->basic_blocks[edge->dst]);
		edge = edge->next;
	}

	dr_fprintf(outfile, "}\n");
}
